package cn.com.basic.framework.aliyun.odps.tunnel;

import cn.com.basic.framework.aliyun.odps.OdpsManager;
import cn.com.basic.framework.aliyun.odps.utils.OdpsUtil;
import cn.com.basic.framework.aliyun.odps.utils.OdpsWriterProxy;
import cn.com.basic.framework.exception.enums.errors.SystemError;
import cn.com.basic.framework.exception.exceptions.SysException;
import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.lang.Assert;
import com.aliyun.odps.Table;
import com.aliyun.odps.data.Record;
import com.aliyun.odps.tunnel.TableTunnel;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

public class UploadTunnel {

    private static Logger logger = LoggerFactory.getLogger(UploadTunnel.class);
    private OdpsManager manager;
    private String tunnelServer;
    private boolean truncate = false;
    private int blockBufferSizeInMB = 8;
    private List<Long> blocks;

    public UploadTunnel(OdpsManager manager) {
        this.manager = manager;
    }

    public void upload(String projectName, String tableName, List<Record> records) {
        upload(projectName, tableName, records, "");
    }

    public void upload(String projectName, String tableName, List<Record> records, String partition) {
        Assert.isTrue(StringUtils.isNotEmpty(projectName), "projectName is not presented");
        Assert.isTrue(StringUtils.isNotEmpty(tableName), "tableName is not presented");
        Assert.isTrue(CollectionUtil.isNotEmpty(records), "at least one record should be added in records");

        if (StringUtils.isNotBlank(partition))
            partition = OdpsUtil.formatPartition(partition);

        Table table = manager.getTable(projectName, tableName);

        OdpsUtil.preCheckPartition(manager.getOdps(), table, partition, this.truncate);

        OdpsUtil.dealTruncate(manager.getOdps(), table, partition, this.truncate);

        TableTunnel tableTunnel = new TableTunnel(manager.getOdps());
        if (StringUtils.isNoneBlank(tunnelServer)) {
            tableTunnel.setEndpoint(tunnelServer);
        }

        TableTunnel.UploadSession uploadSession = OdpsUtil.createMasterTunnelUpload(tableTunnel, projectName, tableName, partition);

        blocks = new ArrayList<Long>();
        AtomicLong blockId = new AtomicLong(0);

        OdpsWriterProxy proxy;
        try {
            proxy = new OdpsWriterProxy(uploadSession, blockBufferSizeInMB, blockId, false);
        } catch (IOException e) {
            logger.error("init OdpsWriterProxy Error : " + e.getMessage(), e);
            throw SysException.newException(SystemError.SYSTEM_ERROR, e.getMessage());
        }

        for (Record record : records) {
            proxy.writeOneRecord(record, blocks);
        }
        proxy.writeRemainingRecord(blocks);


        logger.info("Slave which uploadId=[{}] begin to commit blocks:[\n{}\n].", uploadSession.getId(),
                StringUtils.join(blocks, ","));
        OdpsUtil.masterCompleteBlocks(uploadSession, blocks.toArray(new Long[0]));
        logger.info("Slave which uploadId=[{}] commit blocks ok.", uploadSession.getId());
    }

    public void setTunnelServer(String tunnelServer) {
        this.tunnelServer = tunnelServer;
    }

    public void setTruncate(boolean truncate) {
        this.truncate = truncate;
    }

    public void setBlockBufferSizeInMB(int blockBufferSizeInMB) {
        if (blockBufferSizeInMB < 8) {
            blockBufferSizeInMB = 8;
        }
        if (blockBufferSizeInMB > 512) {
            blockBufferSizeInMB = 512;
        }
        this.blockBufferSizeInMB = blockBufferSizeInMB;
    }
}
